
package com.ctrip.corp.obt.shard.sql.mybatis.interceptor.inner;

import com.ctrip.corp.obt.generic.utils.CollectionUtils;
import com.ctrip.corp.obt.generic.utils.StringPool;
import com.ctrip.corp.obt.generic.utils.StringUtils;
import com.ctrip.corp.obt.shard.annotation.ShardModel;
import com.ctrip.corp.obt.shard.context.DataShardContext;
import com.ctrip.corp.obt.shard.holder.InnerPluginIgnoreHolder;
import com.ctrip.corp.obt.shard.holder.ShardModelHolder;
import com.ctrip.corp.obt.shard.sql.parser.AbstractSqlParserSupport;
import com.ctrip.corp.obt.shard.sql.parser.table.TableRerouteHandler;
import com.ctrip.corp.obt.shard.utils.ClassUtils;
import com.ctrip.corp.obt.shard.utils.PluginUtils;
import com.ctrip.corp.obt.shard.utils.PropertiesWrapper;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.NoArgsConstructor;
import lombok.ToString;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.update.Update;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.sql.Connection;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.Properties;
import java.util.Set;
import java.util.stream.Collectors;

import static com.ctrip.corp.obt.generic.utils.StringPool.UNDERSCORE;

/**
 * TableReroute SQL InnerInterceptor
 *
 * @author marco.zheng
 * @since 2022-02-10
 */
@Data
@NoArgsConstructor
@AllArgsConstructor
@ToString(callSuper = true)
@EqualsAndHashCode(callSuper = true)
@SuppressWarnings({"rawtypes"})
public class TableRerouteInnerPlugin extends AbstractSqlParserSupport implements InnerPlugin {

    private TableRerouteHandler rerouteHandler;


    @Override
    public void beforeQuery(Executor executor, MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) throws SQLException {
        if (InnerPluginIgnoreHolder.willIgnoreTableReroute(ms.getId())) {
            return;
        }
        PluginUtils.DSBoundSql mpBs = PluginUtils.dsBoundSql(boundSql);
        mpBs.sql(parserSingle(mpBs.sql(), null));
    }

    @Override
    public void beforePrepare(StatementHandler sh, Connection connection, Integer transactionTimeout) {
        PluginUtils.DSStatementHandler dsSh = PluginUtils.dsStatementHandler(sh);
        MappedStatement ms = dsSh.mappedStatement();
        SqlCommandType sct = ms.getSqlCommandType();
        if (sct == SqlCommandType.INSERT || sct == SqlCommandType.UPDATE || sct == SqlCommandType.DELETE) {
            if (InnerPluginIgnoreHolder.willIgnoreTableReroute(ms.getId())) {
                return;
            }
            PluginUtils.DSBoundSql mpBs = dsSh.dsBoundSql();
            mpBs.sql(parserMulti(mpBs.sql(), null));
        }
    }


    /**
     * insert 语句处理
     */
    @Override
    protected void processInsert(Insert insert, int index, String sql, Object obj) {
        Table table = insert.getTable();
        if (this.rerouteHandler.doTableFilter(table.getName())) {
            // 过滤退出执行
            return;
        }
        // 重新路由table
        rerouteTable(table);
    }

    /**
     * update 语句处理
     */
    @Override
    protected void processUpdate(Update update, int index, String sql, Object obj) {
        Table table = update.getTable();
        if (this.rerouteHandler.doTableFilter(table.getName())) {
            // 过滤退出执行
            return;
        }
        // 重新路由table
        rerouteTable(table);
    }

    /**
     * delete 语句处理
     */
    @Override
    protected void processDelete(Delete delete, int index, String sql, Object obj) {
        Table table = delete.getTable();
        if (this.rerouteHandler.doTableFilter(table.getName())) {
            // 过滤退出执行
            return;
        }
        // 重新路由table
        rerouteTable(table);
    }

    /**
     * 处理 PlainSelect
     *
     * @param plainSelect ignore
     */
    @Override
    public void processPlainSelect(PlainSelect plainSelect) {
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            Table fromTable = (Table) fromItem;
            if (!this.rerouteHandler.doTableFilter(fromTable.getName())) {
                // 重新路由table
                rerouteTable(fromTable);
            }
        }
    }


    /**
     * 将原始的table进行reroute，table的后缀为 {@link DataShardContext#getTableShardSource()}
     * @param table
     */
    private void rerouteTable(Table table) {
        // 当配置了ShardModel之后，如果未配置targetTables，则对[当前事务中所有的表]按照配置的tableStrategy进行分片
        // 否则只对[targetTables中的表]进行分片处理
        if (Boolean.TRUE.equals(DataShardContext.isAllowTableReroute())) {
            Set<String> routeTables = ShardModelHolder.getShardModelCache().values().stream()
                    .map(ShardModel::targetTables)
                    .flatMap(targetTables -> Arrays.stream(targetTables.split(StringPool.COMMA)))
                    .filter(StringUtils::isNotBlank)
                    .collect(Collectors.toSet());
            if (null != DataShardContext.getTableShardSource()) {
                if (CollectionUtils.isNotEmpty(routeTables) && !routeTables.contains(DataShardContext.getTableShardSource())) {
                    return;
                }
                table.setName(table.getName() + UNDERSCORE +  DataShardContext.getTableShardSource());
            }
        }
    }

    @Override
    public void setProperties(Properties properties) {
        PropertiesWrapper.newInstance(properties).whenNotBlank("tableRerouteHandler",
                ClassUtils::newInstance, this::setRerouteHandler);
    }

    @Override
    public int order() {
        return -50;
    }
}


